package grpc

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"time"

	"gitee.com/zawei/yggdrasil/pkg/errors"
	"gitee.com/zawei/yggdrasil/pkg/metrics"
	"gitee.com/zawei/yggdrasil/pkg/transport/server"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

func extractAID(ctx context.Context) string {
	if md, ok := metadata.FromIncomingContext(ctx); ok {
		return strings.Join(md.Get("aid"), ",")
	}
	return "unknown"
}

func metricUnaryServerInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	startTime := time.Now()
	resp, err := handler(ctx, req)
	var code = int32(http.StatusOK)
	if err != nil {
		st := errors.FromError(err)
		code = st.HttpCode()
	}

	_ = metrics.HistogramObserve(server.MetricsHandleSec, time.Since(startTime).Seconds(), map[string]string{
		"type": "grpc_unary", "method": info.FullMethod, "peer": extractAID(ctx),
	})
	_ = metrics.CounterAdd(server.MetricsHandleTotal, 1, map[string]string{
		"type": "grpc_unary", "method": info.FullMethod, "peer": extractAID(ctx),
		"code": fmt.Sprintf("%d", code),
	})
	return resp, err
}

func metricStreamServerInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	startTime := time.Now()
	err := handler(srv, ss)
	var code = int32(http.StatusOK)
	if err != nil {
		st := errors.FromError(err)
		code = st.HttpCode()
	}
	_ = metrics.HistogramObserve(server.MetricsHandleSec, time.Since(startTime).Seconds(), map[string]string{
		"type": "grpc_stream", "method": info.FullMethod, "peer": extractAID(ss.Context()),
	})
	_ = metrics.CounterAdd(server.MetricsHandleTotal, 1, map[string]string{
		"type": "grpc_stream", "method": info.FullMethod, "peer": extractAID(ss.Context()),
		"code": fmt.Sprintf("%d", code),
	})
	return err
}
